#include <iostream>
#include <cstdlib>
using namespace std;

#include <purewell/iopoll.hpp>
#include <purewell/socket.hpp>
#include <purewell/socketaddress.hpp>
#include <purewell/semaphore.hpp>
using namespace purewell;

#include <signal.h>

#include <openssl/rsa.h>
#include <openssl/crypto.h>
#include <openssl/x509.h>
#include <openssl/pem.h>
#include <openssl/ssl.h>
#include <openssl/err.h>

PWSemaphore* g_lock(PWSemaphore::create(1));

void
log(const char* fmt, ...)
{
	static PWSemaphore* lock(PWSemaphore::create(1));
	lock->lock();
	va_list lst;
	va_start(lst, fmt);
	vfprintf(stdout, fmt, lst);
	fprintf(stdout, "\n");
	fflush(stdout);
	va_end(lst);
	lock->unlock();
}

class XClient : public PWSocket
{
public:
	const char* getSideName(void) const
	{
		return m_client_side ? "CLIENT" : "SERVER";
	}

public:
	XClient(bool client_side, PWIOPoll* io)
	: m_io(io), m_client_side(client_side), m_ctx(NULL),
		m_ssl(NULL), m_cert(NULL), m_method(NULL),
		m_ssl_done(false)
	{
	}

	virtual ~XClient()
	{
	}

	void onEvent(PWIOPoll* io, int fd, int e)
	{
		if ( m_ssl_done )
		{
			if ( e & PWIO_R )
			{
				onRead();
			}
			else
			{
				onWrite();
			}
		}
		else
		{
			procSSL();
			sleep(1);
		}
	}

	void onRead(void)
	{
	}

	void onWrite(void)
	{
	}

	bool procSSL(void)
	{
		if ( m_ssl_done )
		{
			return true;
		}

		int sslret;

		if ( m_client_side )
		{
			log("%s: SSL_connect", getSideName());
			sslret = SSL_connect(m_ssl);
		}
		else
		{
			log("%s: SSL_accept", getSideName());
			sslret = SSL_accept(m_ssl);
		}

		if ( sslret == 1 )
		{
			log("%s: SSL_done!", getSideName());
			m_cert = SSL_get_peer_certificate(m_ssl);
			m_ssl_done = true;
			return true;
		}

		switch(SSL_get_error(m_ssl, sslret))
		{
			case SSL_ERROR_WANT_WRITE:
			{
				log("%s: want to write", getSideName());
				m_io->modify(this, PWIO_W);
				break;
			}
			case SSL_ERROR_WANT_READ:
			{
				log("%s: want to read", getSideName());
				m_io->modify(this, PWIO_R);
				break;
			}
			default:
			{
				log("%s: failed handshake", getSideName());
				exit(EXIT_FAILURE);
			}
		}

		return true;
	}

	bool initSSL(void)
	{
		setModeBlock(false);

		if ( m_client_side )
		{
			m_method = SSLv3_client_method();
			if ( !m_method ) return false;

			m_ctx = SSL_CTX_new(m_method);
			if ( !m_ctx ) return false;

			m_ssl = SSL_new(m_ctx);
			if ( !m_ssl ) return false;

			log("%s: SSL_set_fd", getSideName());
			SSL_set_fd(m_ssl, getHandle());
		}
		else
		{
			m_method = SSLv3_server_method();
			if ( !m_method ) return false;

			m_ctx = SSL_CTX_new(m_method);
			if ( !m_ctx ) return false;

			SSL_CTX_use_certificate_file(m_ctx,
				"server.crt", SSL_FILETYPE_PEM);
			SSL_CTX_use_PrivateKey_file(m_ctx,
				"server.key", SSL_FILETYPE_PEM);

			if ( !SSL_CTX_check_private_key(m_ctx) )
			{
				log("%s: invalid private key", getSideName());
				return false;
			}

			m_ssl = SSL_new(m_ctx);
			if ( !m_ssl ) return false;

			log("%s: SSL_set_fd", getSideName());
			SSL_set_fd(m_ssl, getHandle());
		}

		showCert();

		return true;
	}

	void showCert(void) const
	{
		if ( m_cert )
		{
			log("%s: version: %lu\n", getSideName(),
				X509_get_version(m_cert));
			X509_NAME_print_ex_fp(stdout,
				X509_get_subject_name(m_cert), 0, 0);
			fprintf(stdout, "\n");
			X509_NAME_print_ex_fp(stdout,
				X509_get_subject_name(m_cert), 0, 0);
			fprintf(stdout, "\n");
		}
		else
		{
			log("%s: no cert", getSideName());
		}
	}

	bool sendE(const char* buf, size_t blen)
	{
		return -1 != SSL_write(m_ssl, buf, blen);
	}

	bool receiveE(char* buf, size_t blen)
	{
		return -1 != SSL_read(m_ssl, buf, blen);
	}

public:
	PWIOPoll*	m_io;
	bool		m_client_side;
	SSL_CTX*	m_ctx;
	SSL*		m_ssl;
	X509*		m_cert;
	SSL_METHOD*	m_method;
	bool		m_ssl_done;
	string		m_rbuf;
	string		m_sbuf;
};

bool
initLibrary(void)
{
	signal(SIGPIPE, SIG_IGN);
	SSL_load_error_strings();
	SSLeay_add_ssl_algorithms();
	return true;
}

int
main_server(int argc, char* argv[])
{
	PWIOPoll* io(PWIOPoll::create("select"));
	PWSocket sock;
	sock.initialize(PWSOCK_FAM_IP, PWSOCK_STREAM, 0);
	PWSocketAddress sa;
	sa.setIP(PWSOCK_FAM_IP, "0.0.0.0", "9989");
	sock.setOptionReuseAddress(true);
	sock.bind(sa);
	fprintf(stderr, "Server bind\n");
	sock.listen(5);
	fprintf(stderr, "Server listen\n");

	g_lock->unlock();

	XClient csock(false, io);
	sock.accept(csock);
	sock.destroy();

	io->add(&csock, PWIO_R|PWIO_W);

	fprintf(stderr, "Server accept\n");

	if ( !csock.initSSL() )
	{
		fprintf(stderr, "Server SSL fail\n");
		return EXIT_FAILURE;
	}

	while ( io->dispatch(1000) );

	const char buf[] = "Hello, world!";
	csock.sendE(buf, sizeof(buf));

	sleep(5);

	return EXIT_SUCCESS;
}

int
main_client(int argc, char* argv[])
{
	PWIOPoll* io(PWIOPoll::create("select"));
	XClient sock(true, io);
	sock.initialize(PWSOCK_FAM_IP, PWSOCK_STREAM, 0);
	PWSocketAddress sa;
	sa.setIP(PWSOCK_FAM_IP, "127.0.0.1", "9989");
	sock.connect(sa);
	fprintf(stderr, "Client connect\n");
	if ( !sock.initSSL() )
	{
		fprintf(stderr, "Server SSL fail\n");
		return EXIT_FAILURE;
	}

	io->add(&sock, PWIO_R|PWIO_W);

	while ( io->dispatch(1000) );

	char buf[1024];
	sock.receiveE(buf, sizeof(buf));

	cout << buf << endl;

	return EXIT_SUCCESS;
}

int
main(int argc, char* argv[])
{
	initLibrary();
	g_lock->lock();
	if ( fork() == 0 )
	{
		g_lock->lock();
		return main_client(argc, argv);
	}

	return main_server(argc, argv);
}

